
# Third-Party Imports
import nltk
import openai
from evaluate import load
from datasets import load_dataset

import requests as req
import numpy as np

from sumy.summarizers.lex_rank import LexRankSummarizer
from sumy.summarizers.text_rank import TextRankSummarizer
from sumy.summarizers.lsa import LsaSummarizer
from sumy.summarizers.luhn import LuhnSummarizer
from sumy.parsers.plaintext import PlaintextParser
from sumy.nlp.tokenizers import Tokenizer

import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.offline as pyo
pyo.init_notebook_mode()

from dotenv import load_dotenv
load_dotenv()

# Standard Imports
import os
import json
from string import punctuation


# Plotting functions
def plot_bar_data(*bars, x=None, title="", x_label="", y_label=""):
    fig = go.Figure(
        layout={
            "title": title,
            "xaxis": {"title": x_label},
            "yaxis": {"title": y_label},
            "barmode": "group"
        }, data=[
            go.Bar(name=f"{bar[0]}", x=x, y=bar[1])
            for bar in bars
        ])
    
    return fig

def create_bar(name, data):
    return (name, data)


# Tokenization
def tokenize(doc, remove_stopwords=True):
    banned = list(punctuation)
    
    if remove_stopwords:
        banned += nltk.corpus.stopwords.words("english")
    
    return [
        w.lower() for w in nltk.word_tokenize(doc)
        if w.lower() not in banned
    ]


# Document Summariser Class
# Implementation of all NLP methods used by LAME for text summarisation in a single class
class DocSummariser():
    def __init__(self):
        self._corpus = dict()
    
    def get_corpus(self):
        return self._corpus

    def load_files(self, corpus):
        self._corpus = corpus
    
    def clear_files(self):
        self._corpus = dict()
    
    def _word_tokenize(self, text):
        banned = list(punctuation) + nltk.corpus.stopwords.words("english")
        return [
            w for w in nltk.word_tokenize(text) 
            if w not in banned
        ]
    
    def _chunk_text(self, text, chunk_len):
        chunks = []
        current_chunk = ""
        sents = nltk.sent_tokenize(text)

        for sent in sents:
            if len(nltk.word_tokenize(current_chunk + f" {sent}")) >= chunk_len:
                chunks.append(current_chunk)
                current_chunk = ""
            else:
                current_chunk += f" {sent}"

        chunks.append(current_chunk)

        return chunks
    
    def summarise(self, method, fnames, summary_size):
        # Build input text
        text = " ".join(self._corpus[name] for name in fnames)
        
        # Choose method and return summary
        if method == "se":
            return self._SE_summary(text, summary_size).strip()
        elif method in ("lexR", "texR", "lsa", "luhn"):
            return self._algo_summary(text, method, summary_size).strip()
        elif method == "bart":
            text_chunks = self._chunk_text(text, 400)
            return " ".join(
                self._BART_summary(chunk, summary_size)
                for chunk in text_chunks
            ).strip()
        elif method == "openai":
            text_chunks = self._chunk_text(text, 500)
            return " ".join(
                self._openai_summary(chunk, summary_size)
                for chunk in text_chunks
            ).strip()
    
    def _SE_summary(self, text, summary_size=0.5):
        # Create word and sentence tokens
        words = self._word_tokenize(text)
        word_set = set(words) # set of all unique words in word tokens
        sents = nltk.sent_tokenize(text)

        # Initialise frequency table for word tokens
        w_freq_table = {w: words.count(w) for w in word_set}

        # Score sentences based on frequency of their words
        sent_scores = {
            sent: sum(
                w_freq_table.get(w, 0) 
                for w in self._word_tokenize(sent)
            )
            for sent in sents
        }

        # Build summary
        multiplier = 2 * (1 - summary_size)

        avg = sum(sent_scores.values()) / len(sent_scores)
        summary = " ".join(sent for sent in sents if sent_scores[sent] >= avg * multiplier)
        return summary
    
    def _algo_summary(self, text, method, summary_size=0.5):
        # Get sentence and summary lengths
        sent_length = len(nltk.sent_tokenize(text))
        summary_len = max(int(summary_size * sent_length), 1)

        # Initialise summariser
        if method == "lexR":
            summariser = LexRankSummarizer()
        elif method == "texR":
            summariser = TextRankSummarizer()
        elif method == "lsa":
            summariser = LsaSummarizer()
        elif method == "luhn":
            summariser = LuhnSummarizer()

        # Initialise parser
        parser = PlaintextParser(text, Tokenizer("english"))

        # Create summary
        summary_sents = summariser(parser.document, summary_len)

        return " ".join(str(s) for s in summary_sents)
    
    def _BART_summary(self, text, summary_size=0.5):
        # Get lengths of original text and summary
        word_len = len(nltk.word_tokenize(text))
        summary_len = int((summary_size * word_len) + 0.5)

        # Get API url and headers
        api_url = "https://api-inference.huggingface.co/models/facebook/bart-large-cnn"
        headers = {
            "Authorization": f"Bearer {os.getenv('HUGGING_FACE_API_KEY')}"
        }

        payload = {
            "inputs": text,
            "parameters": {
                "do_sample": False,
                "max_length": min(round(summary_len + 50, -2), word_len),
                "min_length": max(summary_len - 10, 1),
            }
        }

        data = json.dumps(payload)
        res = req.request("POST", api_url, headers=headers, data=data)
        content = json.loads(res.content.decode("utf-8"))

        if isinstance(content, dict):
            return content.get("error", "Something's wrong") + "\n"
        elif isinstance(content, list):
            return content[0].get("summary_text")
    
    def _openai_summary(self, text, summary_size=0.5):
        word_len = len(nltk.word_tokenize(text))
        summary_len = int((summary_size * word_len) + 0.5)

        openai.api_key = os.getenv("OPENAI_API_KEY")
        prompt=f"Summarize the following text in no more than {summary_len} words:\n\n{text}\n\nSummary:"

        max_tokens = round(summary_len + 50, -2)
        if max_tokens < 1: max_tokens = 50
        res = openai.Completion.create(
            model="text-davinci-003", 
            prompt=prompt, 
            temperature=0,
            max_tokens=max_tokens,
            logprobs=0,
        )

        summary = res.choices[0].text
        return summary


def load_article_data(subset_size=5):
    indices = np.random.randint(0, 13368, (subset_size,))
    articles = load_dataset(
        "cnn_dailymail",
        "3.0.0",
        split="validation", 
    ).select(indices)

    return articles


def summarise_sample(sample, method):
    """
    Run a text summarisation method on a single 
    example from the CNN/DailyMail dataset.
    """
    # Get relevant properties from squad sample
    article = sample["article"]
    highlight = sample["highlights"]
    
    # Initialise doc summariser
    doc_summariser = DocSummariser()
    
    # Build and load corpus for doc searcher
    doc_summariser.load_files({"Doc": article})
    
    # Get predicted text
    summary_text = doc_summariser.summarise(method, ["Doc"], 0.1)
    doc_summariser.clear_files()
    
    return summary_text, highlight

def summarise_samples(art_ds, method):
    """
    Run a text summarisation method on multiple 
    examples from the CNN/DailyMail dataset.
    """
    # Initialse lists for storing prediction and reference objects
    predictions = []
    references = []
    
    # Run method on all samples in dataset
    for sample in art_ds:
        summary, highlight = summarise_sample(sample, method)
        predictions.append(summary)
        references.append(highlight)
    
    return predictions, references

def evaluate_method(art_ds, method, rouge_metric):
    """
    Get the average ROUGE scores of a text summarisation
    method after running it on a subset of the CNN/Dailymail
    dataset.
    """
    # Get prediction and reference objects
    preds, refs = summarise_samples(art_ds, method)
    
    # Get results
    results = rouge_metric.compute(predictions=preds, references=refs)
    return results

def visualise_results(results):
    """
    Take results from the evaluate_method function
    an create bar graphs to visualise them.
    """
    method_labels = {
        "se": "Simple Extractive Summarisation",
        "lexR": "LexRank Algorithm",
        "texR": "TextRank Algorithm",
        "lsa": "Latent Semantic Analysis",
        "luhn": "Luhn's Algorithm",
        "bart": "BART",
        "openai": "OpenAI",
    }
    
    plots = dict()
    
    # Create plot for average scores
    x = [method_labels[r["method"]] for r in results]
    rouge1_bar = create_bar("Average ROUGE-1 Score", [r["avg_rouge1"] for r in results])
    rouge2_bar = create_bar("Average ROUGE-2 Score", [r["avg_rouge2"] for r in results])
    rougeL_bar = create_bar("Average ROUGE-L Score", [r["avg_rougeL"] for r in results])
    rougeLsum_bar = create_bar("Average ROUGE-Lsum Score", [r["avg_rougeLsum"] for r in results])
    
    avg_score_plot = plot_bar_data(
        rouge1_bar, 
        rouge2_bar,
        rougeL_bar,
        rougeLsum_bar,
        x=x, 
        title="Average Scores"
    )
    
    plots["average_score_plot"] = avg_score_plot
    
    # Create plot for EM and F1 scores over multiple trials
    for r in results:
        x = [f"Sample #{i+1}" for i in range(len(r["rouge1_scores"]))]
        rouge1_bar = create_bar("ROUGE-1 Score", r["rouge1_scores"])
        rouge2_bar = create_bar("ROUGE-2 Score", r["rouge2_scores"])
        rougeL_bar = create_bar("ROUGE-L Score", r["rougeL_scores"])
        rougeLsum_bar = create_bar("ROUGE-Lsum Score", r["rougeLsum_scores"])
        
        new_plot = plot_bar_data(
            rouge1_bar, 
            rouge2_bar,
            rougeL_bar,
            rougeLsum_bar,
            x=x, 
            title=f"ROUGE Scores for {method_labels[r['method']]}"
        )
        plots[f"{r['method']}_plot"] = new_plot
    
    return plots


def method_evaluator(methods, num_trials=10, dataset_size=50):
    """
    Evaluate several info extraction methods at once.
    """
    # Initialise results object
    results = [
        {
            "rouge1_scores": [],
            "rouge2_scores": [],
            "rougeL_scores": [],
            "rougeLsum_scores": [],
            "method": m
        }
        for m in methods
    ]
    
    # Load squad evaluator
    rouge_metric = load("rouge")
    
    for t in range(num_trials):
        print(f"Trial #{t+1}")
        arts_ds = load_article_data(dataset_size)
        for i, m in enumerate(methods):
            result = evaluate_method(arts_ds, m, rouge_metric)
            results[i]["rouge1_scores"].append(result.get("rouge1", None))
            results[i]["rouge2_scores"].append(result.get("rouge2", None))
            results[i]["rougeL_scores"].append(result.get("rougeL", None))
            results[i]["rougeLsum_scores"].append(result.get("rougeLsum", None))
    
    for i, _ in enumerate(results):
        results[i]["avg_rouge1"] = np.mean(results[i]["rouge1_scores"])
        results[i]["avg_rouge2"] = np.mean(results[i]["rouge2_scores"])
        results[i]["avg_rougeL"] = np.mean(results[i]["rougeL_scores"])
        results[i]["avg_rougeLsum"] = np.mean(results[i]["rougeLsum_scores"])
    
    return results


# Get results of evaluation of each text summarisation method
results = method_evaluator(["se","lexR","texR", "lsa", "luhn", "bart", "openai"], 10, 10)
results

Trial #1

Found cached dataset cnn_dailymail (/Users/bhekimaenetja/.cache/huggingface/datasets/cnn_dailymail/3.0.0/3.0.0/1b3c71476f6d152c31c1730e83ccb08bcf23e348233f4fcc11e182248e6bf7de)

Trial #2

WARNING:datasets.builder:Found cached dataset cnn_dailymail (/Users/bhekimaenetja/.cache/huggingface/datasets/cnn_dailymail/3.0.0/3.0.0/1b3c71476f6d152c31c1730e83ccb08bcf23e348233f4fcc11e182248e6bf7de)

Trial #3

WARNING:datasets.builder:Found cached dataset cnn_dailymail (/Users/bhekimaenetja/.cache/huggingface/datasets/cnn_dailymail/3.0.0/3.0.0/1b3c71476f6d152c31c1730e83ccb08bcf23e348233f4fcc11e182248e6bf7de)

Trial #4

WARNING:datasets.builder:Found cached dataset cnn_dailymail (/Users/bhekimaenetja/.cache/huggingface/datasets/cnn_dailymail/3.0.0/3.0.0/1b3c71476f6d152c31c1730e83ccb08bcf23e348233f4fcc11e182248e6bf7de)

Trial #5

WARNING:datasets.builder:Found cached dataset cnn_dailymail (/Users/bhekimaenetja/.cache/huggingface/datasets/cnn_dailymail/3.0.0/3.0.0/1b3c71476f6d152c31c1730e83ccb08bcf23e348233f4fcc11e182248e6bf7de)


# Get data visualisations of results
results_plots = visualise_results(results)


# Results for simple extractive summarisation
results_plots["se_plot"]


# Results for LexRank
results_plots["lexR_plot"]


# Results for TextRank
results_plots["texR_plot"]


# Results for latent semantic analysis
results_plots["lsa_plot"]


# Results for Luhn's algorithm
results_plots["luhn_plot"]


# Results for BART
results_plots["bart_plot"]


# Results for OpenAI
results_plots["openai_plot"]


# Average ROUGE scores for all methods
results_plots["average_score_plot"]

Imports¶

Utility Functions¶

Plotting¶

Text Processing¶

DocSummariser Class¶

Loading Data¶

Method Evaluation¶

Helper Functions¶

Method Evaluator¶

Evaluation¶